"""
基于LangGraph的旅游助手工作流
实现多轮对话和状态管理
"""

import logging
from typing import Dict, List, Optional, Any, TypedDict, Annotated
from dataclasses import dataclass

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver

from .travel_agent import TravelAgent, AgentResponse
from .intent_classifier import IntentType
from ..utils.logger import setup_logger


class AgentState(TypedDict):
    """Agent状态定义"""
    messages: List[Dict[str, Any]]  # 对话消息列表
    current_query: str  # 当前查询
    intent: str  # 识别的意图
    tool_used: str  # 使用的工具
    response: str  # 生成的响应
    quality_score: float  # 质量得分
    needs_clarification: bool  # 是否需要澄清
    clarification_message: Optional[str]  # 澄清消息
    conversation_context: Dict[str, Any]  # 对话上下文
    error_message: Optional[str]  # 错误消息


class LangGraphWorkflow:
    """基于LangGraph的旅游助手工作流"""
    
    def __init__(self, travel_agent: TravelAgent):
        """
        初始化工作流
        
        Args:
            travel_agent: 旅游助手Agent实例
        """
        self.logger = setup_logger(__name__)
        self.travel_agent = travel_agent
        
        # 创建检查点存储
        self.memory = MemorySaver()
        
        # 创建状态图
        self.workflow = self._create_workflow()
        
        self.logger.info("LangGraph工作流初始化完成")
    
    def _create_workflow(self) -> StateGraph:
        """创建工作流图"""
        # 创建状态图
        workflow = StateGraph(AgentState)
        
        # 添加节点
        workflow.add_node("intent_classification", self._intent_classification_node)
        workflow.add_node("travel_consultation", self._travel_consultation_node)
        workflow.add_node("route_planning", self._route_planning_node)
        workflow.add_node("quality_check", self._quality_check_node)
        workflow.add_node("response_generation", self._response_generation_node)
        workflow.add_node("clarification", self._clarification_node)
        workflow.add_node("error_handling", self._error_handling_node)
        
        # 设置入口点
        workflow.set_entry_point("intent_classification")
        
        # 添加条件边
        workflow.add_conditional_edges(
            "intent_classification",
            self._route_after_intent,
            {
                "travel_consultation": "travel_consultation",
                "route_planning": "route_planning",
                "clarification": "clarification",
                "error": "error_handling"
            }
        )
        
        workflow.add_conditional_edges(
            "travel_consultation",
            self._route_after_travel_consultation,
            {
                "quality_check": "quality_check",
                "error": "error_handling"
            }
        )
        
        workflow.add_conditional_edges(
            "route_planning",
            self._route_after_route_planning,
            {
                "quality_check": "quality_check",
                "clarification": "clarification",
                "error": "error_handling"
            }
        )
        
        workflow.add_conditional_edges(
            "quality_check",
            self._route_after_quality_check,
            {
                "response_generation": "response_generation",
                "clarification": "clarification",
                "error": "error_handling"
            }
        )
        
        # 添加结束边
        workflow.add_edge("response_generation", END)
        workflow.add_edge("clarification", END)
        workflow.add_edge("error_handling", END)
        
        return workflow.compile(checkpointer=self.memory)
    
    def _intent_classification_node(self, state: AgentState) -> AgentState:
        """意图分类节点"""
        try:
            self.logger.info("执行意图分类")
            
            current_query = state["current_query"]
            intent_result = self.travel_agent.intent_classifier.classify_intent(current_query)
            
            # 更新状态
            state["intent"] = intent_result.intent.value
            state["conversation_context"]["last_intent"] = intent_result.intent.value
            
            self.logger.info(f"意图分类完成: {intent_result.intent.value}")
            return state
            
        except Exception as e:
            self.logger.error(f"意图分类失败: {e}")
            state["error_message"] = f"意图分类失败: {str(e)}"
            return state
    
    def _travel_consultation_node(self, state: AgentState) -> AgentState:
        """旅游咨询节点"""
        try:
            self.logger.info("执行旅游咨询")
            
            current_query = state["current_query"]
            response = self.travel_agent._handle_travel_consultation(
                current_query, 
                type('IntentResult', (), {
                    'intent': IntentType.TRAVEL_CONSULTATION,
                    'extracted_params': {}
                })
            )
            
            
            # 更新状态
            state["response"] = response.content
            state["tool_used"] = response.tool_used
            state["needs_clarification"] = response.needs_clarification
            
            self.logger.info("旅游咨询完成")
            return state
            
        except Exception as e:
            self.logger.error(f"旅游咨询失败: {e}")
            state["error_message"] = f"旅游咨询失败: {str(e)}"
            return state
    
    def _route_planning_node(self, state: AgentState) -> AgentState:
        """路线规划节点"""
        try:
            self.logger.info("执行路线规划")
            
            current_query = state["current_query"]
            intent_result = self.travel_agent.intent_classifier.classify_intent(current_query)
            response = self.travel_agent._handle_route_planning(current_query, intent_result)
            
            # 更新状态
            state["response"] = response.content
            state["tool_used"] = response.tool_used
            state["needs_clarification"] = response.needs_clarification
            state["clarification_message"] = response.clarification_message
            
            self.logger.info("路线规划完成")
            return state
            
        except Exception as e:
            self.logger.error(f"路线规划失败: {e}")
            state["error_message"] = f"路线规划失败: {str(e)}"
            return state
    
    def _quality_check_node(self, state: AgentState) -> AgentState:
        """质量检查节点"""
        try:
            self.logger.info("执行质量检查")
            
            current_query = state["current_query"]
            response = state["response"]
            
            # 执行基础质量检查
            quality_result = self.travel_agent.quality_checker.check_quality(
                current_query, response
            )
            
            # 更新状态
            state["quality_score"] = quality_result.score
            
            
            # 如果质量太差，更新响应
            if self.travel_agent.quality_checker.should_reject_answer(quality_result):
                state["response"] = self.travel_agent.quality_checker.get_rejection_message(quality_result)
                state["needs_clarification"] = True
                state["clarification_message"] = "答案质量不佳，请重新提问"
            
            self.logger.info(f"质量检查完成: 得分={quality_result.score:.2f}")
            return state
            
        except Exception as e:
            self.logger.error(f"质量检查失败: {e}")
            state["error_message"] = f"质量检查失败: {str(e)}"
            return state
    
    def _response_generation_node(self, state: AgentState) -> AgentState:
        """响应生成节点"""
        try:
            self.logger.info("生成最终响应")
            
            # 添加消息到对话历史
            message = {
                "role": "assistant",
                "content": state["response"],
                "intent": state["intent"],
                "tool_used": state["tool_used"],
                "quality_score": state["quality_score"]
            }
            
            state["messages"].append(message)
            
            self.logger.info("响应生成完成")
            return state
            
        except Exception as e:
            self.logger.error(f"响应生成失败: {e}")
            state["error_message"] = f"响应生成失败: {str(e)}"
            return state
    
    def _clarification_node(self, state: AgentState) -> AgentState:
        """澄清节点"""
        try:
            self.logger.info("处理澄清请求")
            
            clarification_message = state.get("clarification_message", "请提供更多信息")
            
            # 添加澄清消息到对话历史
            message = {
                "role": "assistant",
                "content": clarification_message,
                "intent": "clarification",
                "tool_used": "clarification",
                "quality_score": 0.0
            }
            
            state["messages"].append(message)
            
            self.logger.info("澄清处理完成")
            return state
            
        except Exception as e:
            self.logger.error(f"澄清处理失败: {e}")
            state["error_message"] = f"澄清处理失败: {str(e)}"
            return state
    
    def _error_handling_node(self, state: AgentState) -> AgentState:
        """错误处理节点"""
        try:
            self.logger.info("处理错误")
            
            error_message = state.get("error_message", "处理请求时出现未知错误")
            
            # 添加错误消息到对话历史
            message = {
                "role": "assistant",
                "content": f"抱歉，{error_message}",
                "intent": "error",
                "tool_used": "error_handler",
                "quality_score": 0.0
            }
            
            state["messages"].append(message)
            
            self.logger.info("错误处理完成")
            return state
            
        except Exception as e:
            self.logger.error(f"错误处理失败: {e}")
            return state
    
    def _route_after_intent(self, state: AgentState) -> str:
        """意图分类后的路由"""
        intent = state.get("intent", "unknown")
        
        if intent == "travel_consultation":
            return "travel_consultation"
        elif intent == "route_planning":
            return "route_planning"
        elif intent == "unknown":
            return "clarification"
        else:
            return "error"
    
    def _route_after_travel_consultation(self, state: AgentState) -> str:
        """旅游咨询后的路由"""
        if state.get("error_message"):
            return "error"
        else:
            return "quality_check"
    
    def _route_after_route_planning(self, state: AgentState) -> str:
        """路线规划后的路由"""
        if state.get("error_message"):
            return "error"
        elif state.get("needs_clarification"):
            return "clarification"
        else:
            return "quality_check"
    
    def _route_after_quality_check(self, state: AgentState) -> str:
        """质量检查后的路由"""
        if state.get("error_message"):
            return "error"
        elif state.get("needs_clarification"):
            return "clarification"
        else:
            return "response_generation"
    
    def process_query(self, query: str, thread_id: str = "default") -> Dict[str, Any]:
        """
        处理用户查询
        
        Args:
            query: 用户查询
            thread_id: 线程ID（用于多轮对话）
            
        Returns:
            处理结果
        """
        try:
            self.logger.info(f"开始处理查询: {query} (线程: {thread_id})")
            
            # 创建初始状态
            initial_state = {
                "messages": [],
                "current_query": query,
                "intent": "",
                "tool_used": "",
                "response": "",
                "quality_score": 0.0,
                "needs_clarification": False,
                "clarification_message": None,
                "conversation_context": self.travel_agent.conversation_context,
                "error_message": None
            }
            
            # 执行工作流
            config = {"configurable": {"thread_id": thread_id}}
            result = self.workflow.invoke(initial_state, config)
            
            # 更新Agent的对话上下文
            self.travel_agent.conversation_context = result["conversation_context"]
            
            self.logger.info(f"查询处理完成: {thread_id}")
            return result
            
        except Exception as e:
            self.logger.error(f"处理查询失败: {e}")
            return {
                "messages": [{
                    "role": "assistant",
                    "content": f"抱歉，处理您的请求时出现错误: {str(e)}",
                    "intent": "error",
                    "tool_used": "error_handler",
                    "quality_score": 0.0
                }],
                "error_message": str(e)
            }
    
    def get_conversation_history(self, thread_id: str = "default") -> List[Dict[str, Any]]:
        """获取对话历史"""
        try:
            config = {"configurable": {"thread_id": thread_id}}
            state = self.workflow.get_state(config)
            return state.values.get("messages", [])
        except Exception as e:
            self.logger.error(f"获取对话历史失败: {e}")
            return []
    
    def reset_conversation(self, thread_id: str = "default"):
        """重置对话"""
        try:
            self.travel_agent.reset_conversation()
            self.logger.info(f"对话已重置: {thread_id}")
        except Exception as e:
            self.logger.error(f"重置对话失败: {e}")
